package com.sz.admin.ai.controller.embedding;

import com.sz.admin.ai.factory.embeddingFactory.EmbeddingHandlerFactory;
import com.sz.admin.ai.factory.embeddingFactory.Handler.EmbeddingHandler;
import com.sz.admin.ai.factory.embeddingFactory.entity.EmbeddingHandlerRequestDTO;
import com.sz.admin.ai.util.EmbeddingUtil;
import com.sz.admin.ai.util.LogAspect;
import com.sz.core.common.entity.ApiResult;
import com.sz.core.util.StringUtils;
import lombok.AllArgsConstructor;
import lombok.extern.log4j.Log4j2;
import org.springframework.ai.document.Document;
import org.springframework.web.bind.annotation.*;

import java.util.List;
import java.util.Map;

/**
 * @描述:
 * @Author: TangYuan
 * @Date: 2025/3/3 10:31
 */
@RequestMapping("/ai/embedding")
@RestController
@CrossOrigin
@AllArgsConstructor
@Log4j2
public class EmbeddingController {
    
    private final EmbeddingHandlerFactory embeddingHandlerFactory;
    
    @LogAspect
    @PostMapping("/doVectorization/fileUrl")
    public ApiResult<Object> fileUrlEmbedding(@RequestBody Map<String, Object> map) {
        try {
            String username = (String) map.get("username");
            String password = (String) map.get("password");
            
            String fileUrl = (String) map.get("fileUrl");
            String fileId = (String) map.get("fileId");
            String fileName = (String) map.get("fileName");
            String platformName = (String) map.get("platformName");
            String modelName = (String) map.get("modelName");
            String userId = (String) map.get("userId");
            String knowledgeBaseId = (String) map.get("knowledgeBaseId");
            
            //? 参数校验   ---    带优化 （现只有非空判断）
            if (StringUtils.isEmpty(fileUrl)) return ApiResult.error("文件地址不能为空");
            if (StringUtils.isEmpty(fileName)) return ApiResult.error("文件名不能为空");
            if (StringUtils.isEmpty(platformName)) return ApiResult.error("模型平台参数错误， 目前支持（ollama，aliyun）");
            if (StringUtils.isEmpty(modelName)) return ApiResult.error("模型名错误");
            if (StringUtils.isEmpty(userId)) return ApiResult.error("用户标识不能为空");
            if (StringUtils.isEmpty(knowledgeBaseId)) return ApiResult.error("知识库标识不能为空");
   
            //? 验证key    ---    待优化
            // if (StringUtils.isEmpty(username) || StringUtils.isEmpty(password)) return ApiResult.error("用户名密码不能为空");
            // ZsjAiUser zsjAiUser = zsjAiUserService.selectZsjAiUserByUsername(username);
            // if (zsjAiUser == null) return ApiResult.error("暂无用户");
            // if (!password.equals(MD5.hash(zsjAiUser.getUsername() + zsjAiUser.getPassword()))) return ApiResult.error("密码不正确");
            // if (zsjAiUser.getEnableFlag() != 1) return ApiResult.error("用户未启用");
            
            // 读取和分割文本
            List<Document> documentList = EmbeddingUtil.readAndSplit(fileUrl);
            // 组合请求体
            EmbeddingHandlerRequestDTO requestDTO = EmbeddingHandlerRequestDTO.builder()
                    .documentList(documentList)
                    .thirdPartyName("11111")
                    .modelName(modelName)
                    .modelKey("sk-f5bee87ddd3347f18faca5e8f6b10b73")
                    .fileId(fileId)
                    .fileName(fileName)
                    .userId(userId)
                    .knowledgeBaseId(knowledgeBaseId)
                    .build();
            
            // 获取模型的Handler
            EmbeddingHandler embeddingHandler = embeddingHandlerFactory.getHandler(platformName);
            // 进行向量化操作
            Boolean b = embeddingHandler.vectorizedData(requestDTO);
            return ApiResult.success(b);
        } catch (Exception e) {
            log.error(e.getMessage());
            e.printStackTrace();
            return ApiResult.error("");
        }
    }
    
    @LogAspect
    @PostMapping(value = "/delete/knowledgeBase")
    public ApiResult<Object> deleteByKnowledgeBase(@RequestBody Map<String, Object> map) {
        try {
            String username = (String) map.get("username");
            String password = (String) map.get("password");
            String platformName = (String) map.get("platformName");
            String modelName = (String) map.get("modelName");
            String userId = (String) map.get("userId");
            String knowledgeBaseId = (String) map.get("knowledgeBaseId");
            
            //? 参数校验   ---    带优化 （现只有非空判断）
            if (StringUtils.isEmpty(platformName)) return ApiResult.error("模型平台参数错误， 目前支持（ollama，aliyun）");
            if (StringUtils.isEmpty(modelName)) return ApiResult.error("模型名错误");
            if (StringUtils.isEmpty(userId)) return ApiResult.error("用户标识不能为空");
            if (StringUtils.isEmpty(knowledgeBaseId)) return ApiResult.error("知识库标识不能为空");
            
            // //? 验证key    ---    待优化
            // if (StringUtils.isEmpty(username) || StringUtils.isEmpty(password)) return ApiResult.error("用户名密码不能为空");
            // ZsjAiUser zsjAiUser = zsjAiUserService.selectZsjAiUserByUsername(username);
            // if (zsjAiUser == null) return ApiResult.error("暂无用户");
            // if (!password.equals(Md5Utils.hash(zsjAiUser.getUsername() + zsjAiUser.getPassword()))) return ApiResult.error("密码不正确");
            // if (zsjAiUser.getEnableFlag() != 1) return ApiResult.error("用户未启用");
            
            EmbeddingHandlerRequestDTO embeddingHandlerRequestDTO = EmbeddingHandlerRequestDTO.builder()
                    .thirdPartyName("")
                    .modelName(modelName)
                    .userId(userId)
                    .knowledgeBaseId(knowledgeBaseId)
                    .build();
            
            EmbeddingHandler ollamaEmbeddingHandler = embeddingHandlerFactory.getHandler(platformName);
            
            Integer delete = ollamaEmbeddingHandler.delete(embeddingHandlerRequestDTO);
            
            return ApiResult.success(delete);
        } catch (Exception e) {
            log.error(e.getMessage());
            return ApiResult.error("");
        }
    }
    
    @LogAspect
    @PostMapping(value = "/delete/fileName")
    public ApiResult<Object> deleteByFileName(@RequestBody Map<String, Object> map) {
        try {
            String username = (String) map.get("username");
            String password = (String) map.get("password");
            String fileId = (String) map.get("fileId");
            String platformName = (String) map.get("platformName");
            String modelName = (String) map.get("modelName");
            String userId = (String) map.get("userId");
            String knowledgeBaseId = (String) map.get("knowledgeBaseId");
            
            //? 参数校验   ---    带优化 （现只有非空判断）
            if (StringUtils.isEmpty(fileId)) return ApiResult.error("文件不能为空");
            if (StringUtils.isEmpty(platformName)) return ApiResult.error("模型平台参数错误， 目前支持（ollama，aliyun）");
            if (StringUtils.isEmpty(modelName)) return ApiResult.error("模型名错误");
            if (StringUtils.isEmpty(userId)) return ApiResult.error("用户标识不能为空");
            if (StringUtils.isEmpty(knowledgeBaseId)) return ApiResult.error("知识库标识不能为空");
            
            // //? 验证key    ---    待优化
            // if (StringUtils.isEmpty(username) || StringUtils.isEmpty(password)) return ApiResult.error("用户名密码不能为空");
            // ZsjAiUser zsjAiUser = zsjAiUserService.selectZsjAiUserByUsername(username);
            // if (zsjAiUser == null) return ApiResult.error("暂无用户");
            // if (!password.equals(Md5Utils.hash(zsjAiUser.getUsername() + zsjAiUser.getPassword()))) return ApiResult.error("密码不正确");
            // if (zsjAiUser.getEnableFlag() != 1) return ApiResult.error("用户未启用");
            
            EmbeddingHandlerRequestDTO embeddingHandlerRequestDTO = EmbeddingHandlerRequestDTO.builder()
                    .thirdPartyName("")
                    .fileId(fileId)
                    .modelName(modelName)
                    .userId(userId)
                    .knowledgeBaseId(knowledgeBaseId)
                    .build();
            
            EmbeddingHandler ollamaEmbeddingHandler = embeddingHandlerFactory.getHandler(platformName);
            
            Integer delete = ollamaEmbeddingHandler.delete(embeddingHandlerRequestDTO);
            
            return ApiResult.success(delete);
        } catch (Exception e) {
            log.error(e.getMessage());
            return ApiResult.error("");
        }
    }
}
